{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# This notebook must be run under PySpark (2.0.2 +)\n",
    "\n",
    "I had better luck when I restarted the notebook kernel in between different parquet groups."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Convert and repartition Citibike DataFrame using Spark"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# standard imports \n",
    "funcs = pyspark.sql.functions\n",
    "types = pyspark.sql.types"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "sqlContext.sql(\"set spark.sql.shuffle.partitions=32\")\n",
    "bike = spark.read.parquet('/data/citibike.parquet')\n",
    "bike.registerTempTable('bike')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style>\n",
       "    .dataframe thead tr:only-child th {\n",
       "        text-align: right;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: left;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>trip_duration</th>\n",
       "      <th>start_time</th>\n",
       "      <th>stop_time</th>\n",
       "      <th>start_station_id</th>\n",
       "      <th>start_station_name</th>\n",
       "      <th>start_station_latitude</th>\n",
       "      <th>start_station_longitude</th>\n",
       "      <th>end_station_id</th>\n",
       "      <th>end_station_name</th>\n",
       "      <th>end_station_latitude</th>\n",
       "      <th>end_station_longitude</th>\n",
       "      <th>bike_id</th>\n",
       "      <th>user_type</th>\n",
       "      <th>birth_year</th>\n",
       "      <th>gender</th>\n",
       "      <th>start_taxizone_id</th>\n",
       "      <th>end_taxizone_id</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>801</td>\n",
       "      <td>1481563027000000</td>\n",
       "      <td>1481563829000000</td>\n",
       "      <td>3002</td>\n",
       "      <td>\"South End Ave &amp; Liberty St\"</td>\n",
       "      <td>40.711512</td>\n",
       "      <td>-74.015756</td>\n",
       "      <td>346</td>\n",
       "      <td>\"Bank St &amp; Hudson St\"</td>\n",
       "      <td>40.736529</td>\n",
       "      <td>-74.006180</td>\n",
       "      <td>25788</td>\n",
       "      <td>\"Subscriber\"</td>\n",
       "      <td>1976.0</td>\n",
       "      <td>1</td>\n",
       "      <td>13.0</td>\n",
       "      <td>158.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>767</td>\n",
       "      <td>1481563031000000</td>\n",
       "      <td>1481563799000000</td>\n",
       "      <td>3224</td>\n",
       "      <td>\"W 13 St &amp; Hudson St\"</td>\n",
       "      <td>40.739974</td>\n",
       "      <td>-74.005139</td>\n",
       "      <td>236</td>\n",
       "      <td>\"St Marks Pl &amp; 2 Ave\"</td>\n",
       "      <td>40.728419</td>\n",
       "      <td>-73.987140</td>\n",
       "      <td>25773</td>\n",
       "      <td>\"Subscriber\"</td>\n",
       "      <td>1973.0</td>\n",
       "      <td>1</td>\n",
       "      <td>249.0</td>\n",
       "      <td>79.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>883</td>\n",
       "      <td>1481563034000000</td>\n",
       "      <td>1481563918000000</td>\n",
       "      <td>3263</td>\n",
       "      <td>\"Cooper Square &amp; E 7 St\"</td>\n",
       "      <td>40.729236</td>\n",
       "      <td>-73.990868</td>\n",
       "      <td>127</td>\n",
       "      <td>\"Barrow St &amp; Hudson St\"</td>\n",
       "      <td>40.731724</td>\n",
       "      <td>-74.006744</td>\n",
       "      <td>20572</td>\n",
       "      <td>\"Subscriber\"</td>\n",
       "      <td>1963.0</td>\n",
       "      <td>1</td>\n",
       "      <td>79.0</td>\n",
       "      <td>158.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>427</td>\n",
       "      <td>1481563036000000</td>\n",
       "      <td>1481563464000000</td>\n",
       "      <td>484</td>\n",
       "      <td>\"W 44 St &amp; 5 Ave\"</td>\n",
       "      <td>40.755003</td>\n",
       "      <td>-73.980144</td>\n",
       "      <td>492</td>\n",
       "      <td>\"W 33 St &amp; 7 Ave\"</td>\n",
       "      <td>40.750200</td>\n",
       "      <td>-73.990931</td>\n",
       "      <td>25620</td>\n",
       "      <td>\"Subscriber\"</td>\n",
       "      <td>1985.0</td>\n",
       "      <td>1</td>\n",
       "      <td>161.0</td>\n",
       "      <td>186.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>478</td>\n",
       "      <td>1481563039000000</td>\n",
       "      <td>1481563518000000</td>\n",
       "      <td>519</td>\n",
       "      <td>\"Pershing Square North\"</td>\n",
       "      <td>40.751873</td>\n",
       "      <td>-73.977706</td>\n",
       "      <td>526</td>\n",
       "      <td>\"E 33 St &amp; 5 Ave\"</td>\n",
       "      <td>40.747659</td>\n",
       "      <td>-73.984907</td>\n",
       "      <td>25019</td>\n",
       "      <td>\"Subscriber\"</td>\n",
       "      <td>1974.0</td>\n",
       "      <td>2</td>\n",
       "      <td>170.0</td>\n",
       "      <td>164.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   trip_duration        start_time         stop_time  start_station_id  \\\n",
       "0            801  1481563027000000  1481563829000000              3002   \n",
       "1            767  1481563031000000  1481563799000000              3224   \n",
       "2            883  1481563034000000  1481563918000000              3263   \n",
       "3            427  1481563036000000  1481563464000000               484   \n",
       "4            478  1481563039000000  1481563518000000               519   \n",
       "\n",
       "             start_station_name  start_station_latitude  \\\n",
       "0  \"South End Ave & Liberty St\"               40.711512   \n",
       "1         \"W 13 St & Hudson St\"               40.739974   \n",
       "2      \"Cooper Square & E 7 St\"               40.729236   \n",
       "3             \"W 44 St & 5 Ave\"               40.755003   \n",
       "4       \"Pershing Square North\"               40.751873   \n",
       "\n",
       "   start_station_longitude  end_station_id         end_station_name  \\\n",
       "0               -74.015756             346    \"Bank St & Hudson St\"   \n",
       "1               -74.005139             236    \"St Marks Pl & 2 Ave\"   \n",
       "2               -73.990868             127  \"Barrow St & Hudson St\"   \n",
       "3               -73.980144             492        \"W 33 St & 7 Ave\"   \n",
       "4               -73.977706             526        \"E 33 St & 5 Ave\"   \n",
       "\n",
       "   end_station_latitude  end_station_longitude  bike_id     user_type  \\\n",
       "0             40.736529             -74.006180    25788  \"Subscriber\"   \n",
       "1             40.728419             -73.987140    25773  \"Subscriber\"   \n",
       "2             40.731724             -74.006744    20572  \"Subscriber\"   \n",
       "3             40.750200             -73.990931    25620  \"Subscriber\"   \n",
       "4             40.747659             -73.984907    25019  \"Subscriber\"   \n",
       "\n",
       "   birth_year  gender  start_taxizone_id  end_taxizone_id  \n",
       "0      1976.0       1               13.0            158.0  \n",
       "1      1973.0       1              249.0             79.0  \n",
       "2      1963.0       1               79.0            158.0  \n",
       "3      1985.0       1              161.0            186.0  \n",
       "4      1974.0       2              170.0            164.0  "
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "spark.sql('select * from bike limit 5').toPandas()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "bike = (bike\n",
    "    .withColumn('start_time', \n",
    "                    funcs.from_unixtime((bike.start_time/1000000).cast(types.IntegerType()))\n",
    "                    .cast(types.TimestampType()))\n",
    "    .withColumn('stop_time', \n",
    "                    funcs.from_unixtime((bike.stop_time/1000000).cast(types.IntegerType()))\n",
    "                    .cast(types.TimestampType()))        \n",
    "#     .withColumn('start_time', \n",
    "#                 (funcs.substring(bike.start_time, 1, 20)).cast(types.TimestampType()))\n",
    "#     .withColumn('stop_time', bike.stop_time.cast(types.TimestampType())) \\\n",
    "    .withColumn('start_taxizone_id', bike.start_taxizone_id.cast(types.FloatType())) \n",
    "    .withColumn('end_taxizone_id', bike.end_taxizone_id.cast(types.FloatType())) \n",
    "        )\n",
    "bike.registerTempTable('bike2')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style>\n",
       "    .dataframe thead tr:only-child th {\n",
       "        text-align: right;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: left;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>trip_duration</th>\n",
       "      <th>start_time</th>\n",
       "      <th>stop_time</th>\n",
       "      <th>start_station_id</th>\n",
       "      <th>start_station_name</th>\n",
       "      <th>start_station_latitude</th>\n",
       "      <th>start_station_longitude</th>\n",
       "      <th>end_station_id</th>\n",
       "      <th>end_station_name</th>\n",
       "      <th>end_station_latitude</th>\n",
       "      <th>end_station_longitude</th>\n",
       "      <th>bike_id</th>\n",
       "      <th>user_type</th>\n",
       "      <th>birth_year</th>\n",
       "      <th>gender</th>\n",
       "      <th>start_taxizone_id</th>\n",
       "      <th>end_taxizone_id</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>801</td>\n",
       "      <td>2016-12-12 12:17:07</td>\n",
       "      <td>2016-12-12 12:30:29</td>\n",
       "      <td>3002</td>\n",
       "      <td>\"South End Ave &amp; Liberty St\"</td>\n",
       "      <td>40.711512</td>\n",
       "      <td>-74.015756</td>\n",
       "      <td>346</td>\n",
       "      <td>\"Bank St &amp; Hudson St\"</td>\n",
       "      <td>40.736529</td>\n",
       "      <td>-74.006180</td>\n",
       "      <td>25788</td>\n",
       "      <td>\"Subscriber\"</td>\n",
       "      <td>1976.0</td>\n",
       "      <td>1</td>\n",
       "      <td>13.0</td>\n",
       "      <td>158.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>767</td>\n",
       "      <td>2016-12-12 12:17:11</td>\n",
       "      <td>2016-12-12 12:29:59</td>\n",
       "      <td>3224</td>\n",
       "      <td>\"W 13 St &amp; Hudson St\"</td>\n",
       "      <td>40.739974</td>\n",
       "      <td>-74.005139</td>\n",
       "      <td>236</td>\n",
       "      <td>\"St Marks Pl &amp; 2 Ave\"</td>\n",
       "      <td>40.728419</td>\n",
       "      <td>-73.987140</td>\n",
       "      <td>25773</td>\n",
       "      <td>\"Subscriber\"</td>\n",
       "      <td>1973.0</td>\n",
       "      <td>1</td>\n",
       "      <td>249.0</td>\n",
       "      <td>79.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>883</td>\n",
       "      <td>2016-12-12 12:17:14</td>\n",
       "      <td>2016-12-12 12:31:58</td>\n",
       "      <td>3263</td>\n",
       "      <td>\"Cooper Square &amp; E 7 St\"</td>\n",
       "      <td>40.729236</td>\n",
       "      <td>-73.990868</td>\n",
       "      <td>127</td>\n",
       "      <td>\"Barrow St &amp; Hudson St\"</td>\n",
       "      <td>40.731724</td>\n",
       "      <td>-74.006744</td>\n",
       "      <td>20572</td>\n",
       "      <td>\"Subscriber\"</td>\n",
       "      <td>1963.0</td>\n",
       "      <td>1</td>\n",
       "      <td>79.0</td>\n",
       "      <td>158.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>427</td>\n",
       "      <td>2016-12-12 12:17:16</td>\n",
       "      <td>2016-12-12 12:24:24</td>\n",
       "      <td>484</td>\n",
       "      <td>\"W 44 St &amp; 5 Ave\"</td>\n",
       "      <td>40.755003</td>\n",
       "      <td>-73.980144</td>\n",
       "      <td>492</td>\n",
       "      <td>\"W 33 St &amp; 7 Ave\"</td>\n",
       "      <td>40.750200</td>\n",
       "      <td>-73.990931</td>\n",
       "      <td>25620</td>\n",
       "      <td>\"Subscriber\"</td>\n",
       "      <td>1985.0</td>\n",
       "      <td>1</td>\n",
       "      <td>161.0</td>\n",
       "      <td>186.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>478</td>\n",
       "      <td>2016-12-12 12:17:19</td>\n",
       "      <td>2016-12-12 12:25:18</td>\n",
       "      <td>519</td>\n",
       "      <td>\"Pershing Square North\"</td>\n",
       "      <td>40.751873</td>\n",
       "      <td>-73.977706</td>\n",
       "      <td>526</td>\n",
       "      <td>\"E 33 St &amp; 5 Ave\"</td>\n",
       "      <td>40.747659</td>\n",
       "      <td>-73.984907</td>\n",
       "      <td>25019</td>\n",
       "      <td>\"Subscriber\"</td>\n",
       "      <td>1974.0</td>\n",
       "      <td>2</td>\n",
       "      <td>170.0</td>\n",
       "      <td>164.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   trip_duration          start_time           stop_time  start_station_id  \\\n",
       "0            801 2016-12-12 12:17:07 2016-12-12 12:30:29              3002   \n",
       "1            767 2016-12-12 12:17:11 2016-12-12 12:29:59              3224   \n",
       "2            883 2016-12-12 12:17:14 2016-12-12 12:31:58              3263   \n",
       "3            427 2016-12-12 12:17:16 2016-12-12 12:24:24               484   \n",
       "4            478 2016-12-12 12:17:19 2016-12-12 12:25:18               519   \n",
       "\n",
       "             start_station_name  start_station_latitude  \\\n",
       "0  \"South End Ave & Liberty St\"               40.711512   \n",
       "1         \"W 13 St & Hudson St\"               40.739974   \n",
       "2      \"Cooper Square & E 7 St\"               40.729236   \n",
       "3             \"W 44 St & 5 Ave\"               40.755003   \n",
       "4       \"Pershing Square North\"               40.751873   \n",
       "\n",
       "   start_station_longitude  end_station_id         end_station_name  \\\n",
       "0               -74.015756             346    \"Bank St & Hudson St\"   \n",
       "1               -74.005139             236    \"St Marks Pl & 2 Ave\"   \n",
       "2               -73.990868             127  \"Barrow St & Hudson St\"   \n",
       "3               -73.980144             492        \"W 33 St & 7 Ave\"   \n",
       "4               -73.977706             526        \"E 33 St & 5 Ave\"   \n",
       "\n",
       "   end_station_latitude  end_station_longitude  bike_id     user_type  \\\n",
       "0             40.736529             -74.006180    25788  \"Subscriber\"   \n",
       "1             40.728419             -73.987140    25773  \"Subscriber\"   \n",
       "2             40.731724             -74.006744    20572  \"Subscriber\"   \n",
       "3             40.750200             -73.990931    25620  \"Subscriber\"   \n",
       "4             40.747659             -73.984907    25019  \"Subscriber\"   \n",
       "\n",
       "   birth_year  gender  start_taxizone_id  end_taxizone_id  \n",
       "0      1976.0       1               13.0            158.0  \n",
       "1      1973.0       1              249.0             79.0  \n",
       "2      1963.0       1               79.0            158.0  \n",
       "3      1985.0       1              161.0            186.0  \n",
       "4      1974.0       2              170.0            164.0  "
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "spark.sql('select * from bike2 limit 5').toPandas()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "bike.sort('start_time') \\\n",
    "    .write.parquet('/data/citibike_spark.parquet', compression='snappy', mode='overwrite')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "# Convert and repartition Subway Dataframe using PySpark"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# standard imports \n",
    "funcs = pyspark.sql.functions\n",
    "types = pyspark.sql.types"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "sqlContext.sql(\"set spark.sql.shuffle.partitions=32\")\n",
    "subway = spark.read.parquet('/data/subway.parquet')\n",
    "subway.registerTempTable(\"subway\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style>\n",
       "    .dataframe thead tr:only-child th {\n",
       "        text-align: right;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: left;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>ca</th>\n",
       "      <th>unit</th>\n",
       "      <th>scp</th>\n",
       "      <th>station</th>\n",
       "      <th>linename</th>\n",
       "      <th>division</th>\n",
       "      <th>endtime</th>\n",
       "      <th>description</th>\n",
       "      <th>cumul_entries</th>\n",
       "      <th>cumul_exits</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>\"A002\"</td>\n",
       "      <td>\"R051\"</td>\n",
       "      <td>\"02-00-00\"</td>\n",
       "      <td>\"NULL\"</td>\n",
       "      <td>\"NULL\"</td>\n",
       "      <td>\"NULL\"</td>\n",
       "      <td>1356145200000000</td>\n",
       "      <td>\"REGULAR\"</td>\n",
       "      <td>3922274</td>\n",
       "      <td>1352720</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>\"A002\"</td>\n",
       "      <td>\"R051\"</td>\n",
       "      <td>\"02-00-00\"</td>\n",
       "      <td>\"NULL\"</td>\n",
       "      <td>\"NULL\"</td>\n",
       "      <td>\"NULL\"</td>\n",
       "      <td>1356159600000000</td>\n",
       "      <td>\"REGULAR\"</td>\n",
       "      <td>3922288</td>\n",
       "      <td>1352730</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>\"A002\"</td>\n",
       "      <td>\"R051\"</td>\n",
       "      <td>\"02-00-00\"</td>\n",
       "      <td>\"NULL\"</td>\n",
       "      <td>\"NULL\"</td>\n",
       "      <td>\"NULL\"</td>\n",
       "      <td>1356174000000000</td>\n",
       "      <td>\"REGULAR\"</td>\n",
       "      <td>3922352</td>\n",
       "      <td>1352824</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>\"A002\"</td>\n",
       "      <td>\"R051\"</td>\n",
       "      <td>\"02-00-00\"</td>\n",
       "      <td>\"NULL\"</td>\n",
       "      <td>\"NULL\"</td>\n",
       "      <td>\"NULL\"</td>\n",
       "      <td>1356188400000000</td>\n",
       "      <td>\"REGULAR\"</td>\n",
       "      <td>3922544</td>\n",
       "      <td>1352881</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>\"A002\"</td>\n",
       "      <td>\"R051\"</td>\n",
       "      <td>\"02-00-00\"</td>\n",
       "      <td>\"NULL\"</td>\n",
       "      <td>\"NULL\"</td>\n",
       "      <td>\"NULL\"</td>\n",
       "      <td>1356202800000000</td>\n",
       "      <td>\"REGULAR\"</td>\n",
       "      <td>3922944</td>\n",
       "      <td>1352967</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "       ca    unit         scp station linename division           endtime  \\\n",
       "0  \"A002\"  \"R051\"  \"02-00-00\"  \"NULL\"   \"NULL\"   \"NULL\"  1356145200000000   \n",
       "1  \"A002\"  \"R051\"  \"02-00-00\"  \"NULL\"   \"NULL\"   \"NULL\"  1356159600000000   \n",
       "2  \"A002\"  \"R051\"  \"02-00-00\"  \"NULL\"   \"NULL\"   \"NULL\"  1356174000000000   \n",
       "3  \"A002\"  \"R051\"  \"02-00-00\"  \"NULL\"   \"NULL\"   \"NULL\"  1356188400000000   \n",
       "4  \"A002\"  \"R051\"  \"02-00-00\"  \"NULL\"   \"NULL\"   \"NULL\"  1356202800000000   \n",
       "\n",
       "  description  cumul_entries  cumul_exits  \n",
       "0   \"REGULAR\"        3922274      1352720  \n",
       "1   \"REGULAR\"        3922288      1352730  \n",
       "2   \"REGULAR\"        3922352      1352824  \n",
       "3   \"REGULAR\"        3922544      1352881  \n",
       "4   \"REGULAR\"        3922944      1352967  "
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "spark.sql('select * from subway limit 5').toPandas()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "subway = \\\n",
    "    subway.withColumn('endtime', funcs.from_unixtime((subway.endtime/1000000).cast(types.IntegerType())) \\\n",
    "                      .cast(types.TimestampType()))\n",
    "subway.registerTempTable(\"subway2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style>\n",
       "    .dataframe thead tr:only-child th {\n",
       "        text-align: right;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: left;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>ca</th>\n",
       "      <th>unit</th>\n",
       "      <th>scp</th>\n",
       "      <th>station</th>\n",
       "      <th>linename</th>\n",
       "      <th>division</th>\n",
       "      <th>endtime</th>\n",
       "      <th>description</th>\n",
       "      <th>cumul_entries</th>\n",
       "      <th>cumul_exits</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>\"A002\"</td>\n",
       "      <td>\"R051\"</td>\n",
       "      <td>\"02-00-00\"</td>\n",
       "      <td>\"NULL\"</td>\n",
       "      <td>\"NULL\"</td>\n",
       "      <td>\"NULL\"</td>\n",
       "      <td>2012-12-21 22:00:00</td>\n",
       "      <td>\"REGULAR\"</td>\n",
       "      <td>3922274</td>\n",
       "      <td>1352720</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>\"A002\"</td>\n",
       "      <td>\"R051\"</td>\n",
       "      <td>\"02-00-00\"</td>\n",
       "      <td>\"NULL\"</td>\n",
       "      <td>\"NULL\"</td>\n",
       "      <td>\"NULL\"</td>\n",
       "      <td>2012-12-22 02:00:00</td>\n",
       "      <td>\"REGULAR\"</td>\n",
       "      <td>3922288</td>\n",
       "      <td>1352730</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>\"A002\"</td>\n",
       "      <td>\"R051\"</td>\n",
       "      <td>\"02-00-00\"</td>\n",
       "      <td>\"NULL\"</td>\n",
       "      <td>\"NULL\"</td>\n",
       "      <td>\"NULL\"</td>\n",
       "      <td>2012-12-22 06:00:00</td>\n",
       "      <td>\"REGULAR\"</td>\n",
       "      <td>3922352</td>\n",
       "      <td>1352824</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>\"A002\"</td>\n",
       "      <td>\"R051\"</td>\n",
       "      <td>\"02-00-00\"</td>\n",
       "      <td>\"NULL\"</td>\n",
       "      <td>\"NULL\"</td>\n",
       "      <td>\"NULL\"</td>\n",
       "      <td>2012-12-22 10:00:00</td>\n",
       "      <td>\"REGULAR\"</td>\n",
       "      <td>3922544</td>\n",
       "      <td>1352881</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>\"A002\"</td>\n",
       "      <td>\"R051\"</td>\n",
       "      <td>\"02-00-00\"</td>\n",
       "      <td>\"NULL\"</td>\n",
       "      <td>\"NULL\"</td>\n",
       "      <td>\"NULL\"</td>\n",
       "      <td>2012-12-22 14:00:00</td>\n",
       "      <td>\"REGULAR\"</td>\n",
       "      <td>3922944</td>\n",
       "      <td>1352967</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "       ca    unit         scp station linename division             endtime  \\\n",
       "0  \"A002\"  \"R051\"  \"02-00-00\"  \"NULL\"   \"NULL\"   \"NULL\" 2012-12-21 22:00:00   \n",
       "1  \"A002\"  \"R051\"  \"02-00-00\"  \"NULL\"   \"NULL\"   \"NULL\" 2012-12-22 02:00:00   \n",
       "2  \"A002\"  \"R051\"  \"02-00-00\"  \"NULL\"   \"NULL\"   \"NULL\" 2012-12-22 06:00:00   \n",
       "3  \"A002\"  \"R051\"  \"02-00-00\"  \"NULL\"   \"NULL\"   \"NULL\" 2012-12-22 10:00:00   \n",
       "4  \"A002\"  \"R051\"  \"02-00-00\"  \"NULL\"   \"NULL\"   \"NULL\" 2012-12-22 14:00:00   \n",
       "\n",
       "  description  cumul_entries  cumul_exits  \n",
       "0   \"REGULAR\"        3922274      1352720  \n",
       "1   \"REGULAR\"        3922288      1352730  \n",
       "2   \"REGULAR\"        3922352      1352824  \n",
       "3   \"REGULAR\"        3922544      1352881  \n",
       "4   \"REGULAR\"        3922944      1352967  "
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "spark.sql('select * from subway2 limit 5').toPandas()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "subway = subway.sort(\"ca\", \"unit\", \"scp\", \"endtime\")\n",
    "subway.write.parquet('/data/subway_spark.parquet', compression='snappy', mode='overwrite')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "# Convert, repartition, and sort Taxi Dataframe using PySpark"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# standard imports \n",
    "funcs = pyspark.sql.functions\n",
    "types = pyspark.sql.types"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "sqlContext.sql(\"set spark.sql.shuffle.partitions=800\")\n",
    "taxi = spark.read.parquet('/data/all_trips_unprocessed.parquet')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "taxi = (\n",
    "    taxi.withColumn('dropoff_datetime', \n",
    "                    funcs.from_unixtime((taxi.dropoff_datetime/1000000).cast(types.IntegerType()))\n",
    "                    .cast(types.TimestampType())) \\\n",
    "    .withColumn('pickup_datetime', \n",
    "                    funcs.from_unixtime((taxi.pickup_datetime/1000000).cast(types.IntegerType()))\n",
    "                .cast(types.TimestampType())) \\\n",
    "    .withColumn('dropoff_taxizone_id', taxi.dropoff_taxizone_id.cast(types.IntegerType())) \\\n",
    "    .withColumn('pickup_taxizone_id', taxi.pickup_taxizone_id.cast(types.IntegerType())) \\\n",
    "    .withColumn('dropoff_latitude', taxi.dropoff_latitude.cast(types.FloatType())) \\\n",
    "    .withColumn('dropoff_longitude', taxi.dropoff_longitude.cast(types.FloatType())) \\\n",
    "    .withColumn('ehail_fee', taxi.ehail_fee.cast(types.FloatType())) \\\n",
    "    .withColumn('extra', taxi.extra.cast(types.FloatType())) \\\n",
    "    .withColumn('fare_amount', taxi.fare_amount.cast(types.FloatType())) \\\n",
    "    .withColumn('improvement_surcharge', taxi.improvement_surcharge.cast(types.FloatType())) \\\n",
    "    .withColumn('mta_tax', taxi.mta_tax.cast(types.FloatType())) \\\n",
    "    .withColumn('pickup_latitude', taxi.pickup_latitude.cast(types.FloatType())) \\\n",
    "    .withColumn('pickup_longitude', taxi.pickup_longitude.cast(types.FloatType())) \\\n",
    "    .withColumn('tip_amount', taxi.tip_amount.cast(types.FloatType())) \\\n",
    "    .withColumn('tolls_amount', taxi.tolls_amount.cast(types.FloatType())) \\\n",
    "    .withColumn('total_amount', taxi.total_amount.cast(types.FloatType())) \\\n",
    "    .withColumn('trip_distance', taxi.trip_distance.cast(types.FloatType())) \\\n",
    "    .withColumn('passenger_count', taxi.passenger_count.cast(types.IntegerType())) \\\n",
    "    .withColumn('rate_code_id', taxi.rate_code_id.cast(types.IntegerType())) \n",
    "#     .withColumn('trip_id', funcs.monotonically_increasing_id())\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "taxi.sort('pickup_datetime').withColumn('trip_id', funcs.monotonically_increasing_id()) \\\n",
    "    .write.parquet('/data/all_trips_spark.parquet', compression='snappy', mode='overwrite')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "# Using Dask, Read, and set index on Taxi Dataframe produced using PySpark, then write to disk for easy reading in Dask\n",
    "\n",
    "For some reason I don't yet understand, this code seems dependent on Dask 0.14.3, and breaks in 0.15."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "pd.options.display.max_rows = 100\n",
    "pd.options.display.max_columns = 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import dask.dataframe as dd\n",
    "import dask.distributed \n",
    "client = dask.distributed.Client()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "trips = dd.read_parquet('/data/all_trips_spark.parquet', engine='arrow')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style>\n",
       "    .dataframe thead tr:only-child th {\n",
       "        text-align: right;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: left;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>dropoff_datetime</th>\n",
       "      <th>dropoff_latitude</th>\n",
       "      <th>dropoff_longitude</th>\n",
       "      <th>dropoff_taxizone_id</th>\n",
       "      <th>ehail_fee</th>\n",
       "      <th>extra</th>\n",
       "      <th>fare_amount</th>\n",
       "      <th>improvement_surcharge</th>\n",
       "      <th>mta_tax</th>\n",
       "      <th>passenger_count</th>\n",
       "      <th>payment_type</th>\n",
       "      <th>pickup_datetime</th>\n",
       "      <th>pickup_latitude</th>\n",
       "      <th>pickup_longitude</th>\n",
       "      <th>pickup_taxizone_id</th>\n",
       "      <th>rate_code_id</th>\n",
       "      <th>store_and_fwd_flag</th>\n",
       "      <th>tip_amount</th>\n",
       "      <th>tolls_amount</th>\n",
       "      <th>total_amount</th>\n",
       "      <th>trip_distance</th>\n",
       "      <th>trip_type</th>\n",
       "      <th>vendor_id</th>\n",
       "      <th>trip_id</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>2009-01-01 00:04:12</td>\n",
       "      <td>40.777058</td>\n",
       "      <td>-73.949608</td>\n",
       "      <td>263.0</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.0</td>\n",
       "      <td>5.800000</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>\"Cash\"</td>\n",
       "      <td>2009-01-01 00:00:00</td>\n",
       "      <td>40.771244</td>\n",
       "      <td>-73.965919</td>\n",
       "      <td>237.0</td>\n",
       "      <td>NaN</td>\n",
       "      <td>None</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>5.800000</td>\n",
       "      <td>1.3</td>\n",
       "      <td>\"yellow\"</td>\n",
       "      <td>\"CMT\"</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2009-01-01 00:05:03</td>\n",
       "      <td>40.735703</td>\n",
       "      <td>-74.005936</td>\n",
       "      <td>249.0</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.0</td>\n",
       "      <td>5.400000</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>\"Cash\"</td>\n",
       "      <td>2009-01-01 00:00:00</td>\n",
       "      <td>40.725952</td>\n",
       "      <td>-73.997482</td>\n",
       "      <td>114.0</td>\n",
       "      <td>NaN</td>\n",
       "      <td>None</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>5.400000</td>\n",
       "      <td>0.9</td>\n",
       "      <td>\"yellow\"</td>\n",
       "      <td>\"CMT\"</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>2009-01-01 00:05:40</td>\n",
       "      <td>40.773746</td>\n",
       "      <td>-73.977753</td>\n",
       "      <td>43.0</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.0</td>\n",
       "      <td>5.800000</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>\"Cash\"</td>\n",
       "      <td>2009-01-01 00:00:02</td>\n",
       "      <td>40.767391</td>\n",
       "      <td>-73.964798</td>\n",
       "      <td>237.0</td>\n",
       "      <td>NaN</td>\n",
       "      <td>None</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>5.800000</td>\n",
       "      <td>1.0</td>\n",
       "      <td>\"yellow\"</td>\n",
       "      <td>\"CMT\"</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>2009-01-01 00:03:08</td>\n",
       "      <td>40.709358</td>\n",
       "      <td>-74.013466</td>\n",
       "      <td>261.0</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.0</td>\n",
       "      <td>4.600000</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>\"Cash\"</td>\n",
       "      <td>2009-01-01 00:00:04</td>\n",
       "      <td>40.708832</td>\n",
       "      <td>-74.011597</td>\n",
       "      <td>261.0</td>\n",
       "      <td>NaN</td>\n",
       "      <td>None</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>4.600000</td>\n",
       "      <td>0.8</td>\n",
       "      <td>\"yellow\"</td>\n",
       "      <td>\"CMT\"</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>2009-01-01 00:19:01</td>\n",
       "      <td>40.712368</td>\n",
       "      <td>-73.944580</td>\n",
       "      <td>80.0</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.0</td>\n",
       "      <td>27.799999</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>\"Cash\"</td>\n",
       "      <td>2009-01-01 00:00:07</td>\n",
       "      <td>40.718578</td>\n",
       "      <td>-74.000648</td>\n",
       "      <td>144.0</td>\n",
       "      <td>NaN</td>\n",
       "      <td>None</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>27.799999</td>\n",
       "      <td>5.5</td>\n",
       "      <td>\"yellow\"</td>\n",
       "      <td>\"CMT\"</td>\n",
       "      <td>4</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "     dropoff_datetime  dropoff_latitude  dropoff_longitude  \\\n",
       "0 2009-01-01 00:04:12         40.777058         -73.949608   \n",
       "1 2009-01-01 00:05:03         40.735703         -74.005936   \n",
       "2 2009-01-01 00:05:40         40.773746         -73.977753   \n",
       "3 2009-01-01 00:03:08         40.709358         -74.013466   \n",
       "4 2009-01-01 00:19:01         40.712368         -73.944580   \n",
       "\n",
       "   dropoff_taxizone_id  ehail_fee  extra  fare_amount  improvement_surcharge  \\\n",
       "0                263.0        NaN    0.0     5.800000                    NaN   \n",
       "1                249.0        NaN    0.0     5.400000                    NaN   \n",
       "2                 43.0        NaN    0.0     5.800000                    NaN   \n",
       "3                261.0        NaN    0.0     4.600000                    NaN   \n",
       "4                 80.0        NaN    0.0    27.799999                    NaN   \n",
       "\n",
       "   mta_tax  passenger_count payment_type     pickup_datetime  pickup_latitude  \\\n",
       "0      NaN              NaN       \"Cash\" 2009-01-01 00:00:00        40.771244   \n",
       "1      NaN              NaN       \"Cash\" 2009-01-01 00:00:00        40.725952   \n",
       "2      NaN              NaN       \"Cash\" 2009-01-01 00:00:02        40.767391   \n",
       "3      NaN              NaN       \"Cash\" 2009-01-01 00:00:04        40.708832   \n",
       "4      NaN              NaN       \"Cash\" 2009-01-01 00:00:07        40.718578   \n",
       "\n",
       "   pickup_longitude  pickup_taxizone_id  rate_code_id store_and_fwd_flag  \\\n",
       "0        -73.965919               237.0           NaN               None   \n",
       "1        -73.997482               114.0           NaN               None   \n",
       "2        -73.964798               237.0           NaN               None   \n",
       "3        -74.011597               261.0           NaN               None   \n",
       "4        -74.000648               144.0           NaN               None   \n",
       "\n",
       "   tip_amount  tolls_amount  total_amount  trip_distance trip_type vendor_id  \\\n",
       "0         0.0           0.0      5.800000            1.3  \"yellow\"     \"CMT\"   \n",
       "1         0.0           0.0      5.400000            0.9  \"yellow\"     \"CMT\"   \n",
       "2         0.0           0.0      5.800000            1.0  \"yellow\"     \"CMT\"   \n",
       "3         0.0           0.0      4.600000            0.8  \"yellow\"     \"CMT\"   \n",
       "4         0.0           0.0     27.799999            5.5  \"yellow\"     \"CMT\"   \n",
       "\n",
       "   trip_id  \n",
       "0        0  \n",
       "1        1  \n",
       "2        2  \n",
       "3        3  \n",
       "4        4  "
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trips.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style>\n",
       "    .dataframe thead tr:only-child th {\n",
       "        text-align: right;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: left;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>dropoff_datetime</th>\n",
       "      <th>dropoff_latitude</th>\n",
       "      <th>dropoff_longitude</th>\n",
       "      <th>dropoff_taxizone_id</th>\n",
       "      <th>ehail_fee</th>\n",
       "      <th>extra</th>\n",
       "      <th>fare_amount</th>\n",
       "      <th>improvement_surcharge</th>\n",
       "      <th>mta_tax</th>\n",
       "      <th>passenger_count</th>\n",
       "      <th>payment_type</th>\n",
       "      <th>pickup_datetime</th>\n",
       "      <th>pickup_latitude</th>\n",
       "      <th>pickup_longitude</th>\n",
       "      <th>pickup_taxizone_id</th>\n",
       "      <th>rate_code_id</th>\n",
       "      <th>store_and_fwd_flag</th>\n",
       "      <th>tip_amount</th>\n",
       "      <th>tolls_amount</th>\n",
       "      <th>total_amount</th>\n",
       "      <th>trip_distance</th>\n",
       "      <th>trip_type</th>\n",
       "      <th>vendor_id</th>\n",
       "      <th>trip_id</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>1432504</th>\n",
       "      <td>2017-01-01 00:07:47</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>36</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.5</td>\n",
       "      <td>7.0</td>\n",
       "      <td>0.3</td>\n",
       "      <td>0.5</td>\n",
       "      <td>NaN</td>\n",
       "      <td>\"2\"</td>\n",
       "      <td>2016-12-31 23:59:57</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>36</td>\n",
       "      <td>NaN</td>\n",
       "      <td>\"N\"</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.0</td>\n",
       "      <td>8.300000</td>\n",
       "      <td>1.60</td>\n",
       "      <td>\"green\"</td>\n",
       "      <td>\"2\"</td>\n",
       "      <td>6863359171512</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1432505</th>\n",
       "      <td>2017-01-01 00:15:29</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>63</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.0</td>\n",
       "      <td>16.5</td>\n",
       "      <td>0.3</td>\n",
       "      <td>0.5</td>\n",
       "      <td>NaN</td>\n",
       "      <td>\"2\"</td>\n",
       "      <td>2016-12-31 23:59:58</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>76</td>\n",
       "      <td>NaN</td>\n",
       "      <td>\"N\"</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.0</td>\n",
       "      <td>17.299999</td>\n",
       "      <td>5.20</td>\n",
       "      <td>\"green\"</td>\n",
       "      <td>\"1\"</td>\n",
       "      <td>6863359171513</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1432506</th>\n",
       "      <td>2017-01-01 00:39:07</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>161</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.5</td>\n",
       "      <td>33.5</td>\n",
       "      <td>0.3</td>\n",
       "      <td>0.5</td>\n",
       "      <td>NaN</td>\n",
       "      <td>\"1\"</td>\n",
       "      <td>2016-12-31 23:59:58</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>168</td>\n",
       "      <td>NaN</td>\n",
       "      <td>\"N\"</td>\n",
       "      <td>6.96</td>\n",
       "      <td>0.0</td>\n",
       "      <td>41.759998</td>\n",
       "      <td>8.83</td>\n",
       "      <td>\"green\"</td>\n",
       "      <td>\"2\"</td>\n",
       "      <td>6863359171514</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1432507</th>\n",
       "      <td>2017-01-01 00:03:50</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>209</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.5</td>\n",
       "      <td>5.0</td>\n",
       "      <td>0.3</td>\n",
       "      <td>0.5</td>\n",
       "      <td>NaN</td>\n",
       "      <td>\"2\"</td>\n",
       "      <td>2016-12-31 23:59:58</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>144</td>\n",
       "      <td>NaN</td>\n",
       "      <td>\"N\"</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.0</td>\n",
       "      <td>6.300000</td>\n",
       "      <td>0.70</td>\n",
       "      <td>\"yellow\"</td>\n",
       "      <td>\"1\"</td>\n",
       "      <td>6863359171515</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1432508</th>\n",
       "      <td>2017-01-01 00:14:30</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>134</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.5</td>\n",
       "      <td>13.0</td>\n",
       "      <td>0.3</td>\n",
       "      <td>0.5</td>\n",
       "      <td>NaN</td>\n",
       "      <td>\"2\"</td>\n",
       "      <td>2016-12-31 23:59:59</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>135</td>\n",
       "      <td>NaN</td>\n",
       "      <td>\"N\"</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.0</td>\n",
       "      <td>14.300000</td>\n",
       "      <td>3.41</td>\n",
       "      <td>\"green\"</td>\n",
       "      <td>\"2\"</td>\n",
       "      <td>6863359171516</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "           dropoff_datetime  dropoff_latitude  dropoff_longitude  \\\n",
       "1432504 2017-01-01 00:07:47               NaN                NaN   \n",
       "1432505 2017-01-01 00:15:29               NaN                NaN   \n",
       "1432506 2017-01-01 00:39:07               NaN                NaN   \n",
       "1432507 2017-01-01 00:03:50               NaN                NaN   \n",
       "1432508 2017-01-01 00:14:30               NaN                NaN   \n",
       "\n",
       "         dropoff_taxizone_id  ehail_fee  extra  fare_amount  \\\n",
       "1432504                   36        NaN    0.5          7.0   \n",
       "1432505                   63        NaN    0.0         16.5   \n",
       "1432506                  161        NaN    0.5         33.5   \n",
       "1432507                  209        NaN    0.5          5.0   \n",
       "1432508                  134        NaN    0.5         13.0   \n",
       "\n",
       "         improvement_surcharge  mta_tax  passenger_count payment_type  \\\n",
       "1432504                    0.3      0.5              NaN          \"2\"   \n",
       "1432505                    0.3      0.5              NaN          \"2\"   \n",
       "1432506                    0.3      0.5              NaN          \"1\"   \n",
       "1432507                    0.3      0.5              NaN          \"2\"   \n",
       "1432508                    0.3      0.5              NaN          \"2\"   \n",
       "\n",
       "            pickup_datetime  pickup_latitude  pickup_longitude  \\\n",
       "1432504 2016-12-31 23:59:57              NaN               NaN   \n",
       "1432505 2016-12-31 23:59:58              NaN               NaN   \n",
       "1432506 2016-12-31 23:59:58              NaN               NaN   \n",
       "1432507 2016-12-31 23:59:58              NaN               NaN   \n",
       "1432508 2016-12-31 23:59:59              NaN               NaN   \n",
       "\n",
       "         pickup_taxizone_id  rate_code_id store_and_fwd_flag  tip_amount  \\\n",
       "1432504                  36           NaN                \"N\"        0.00   \n",
       "1432505                  76           NaN                \"N\"        0.00   \n",
       "1432506                 168           NaN                \"N\"        6.96   \n",
       "1432507                 144           NaN                \"N\"        0.00   \n",
       "1432508                 135           NaN                \"N\"        0.00   \n",
       "\n",
       "         tolls_amount  total_amount  trip_distance trip_type vendor_id  \\\n",
       "1432504           0.0      8.300000           1.60   \"green\"       \"2\"   \n",
       "1432505           0.0     17.299999           5.20   \"green\"       \"1\"   \n",
       "1432506           0.0     41.759998           8.83   \"green\"       \"2\"   \n",
       "1432507           0.0      6.300000           0.70  \"yellow\"       \"1\"   \n",
       "1432508           0.0     14.300000           3.41   \"green\"       \"2\"   \n",
       "\n",
       "               trip_id  \n",
       "1432504  6863359171512  \n",
       "1432505  6863359171513  \n",
       "1432506  6863359171514  \n",
       "1432507  6863359171515  \n",
       "1432508  6863359171516  "
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trips.tail()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# arrow engine adds quotes to all string fields for some reason. Strip them out.\n",
    "dtypedict = dict(trips.dtypes)\n",
    "for k in dtypedict:\n",
    "    if dtypedict[k] == np.dtype('O'):\n",
    "        trips[k] = trips[k].str.strip('\"')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": false,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "trips = trips.set_index('pickup_datetime', npartitions=trips.npartitions, sorted=True, compute=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "trips.to_parquet('/data/all_trips.parquet', has_nulls=True, compression=\"SNAPPY\", object_encoding='json')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python [default]",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
